package com.crpc.core.loadBalance.impl;

import com.crpc.core.loadBalance.AbstractLoadBalance;
import com.crpc.core.loadBalance.Selector;
import com.crpc.core.transport.message.CrpcRequest;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

import static com.crpc.core.CrpcBootStrap.CRPC_REQUEST_THREAD_LOCAL;

/**
 * @author wzy
 * @version 1.0
 * @description:
 * @date 2023/10/1 11:52
 */
@Slf4j
public class ConsistentHashLoadBalance extends AbstractLoadBalance {
    @Override
    public Selector getSelector(List<InetSocketAddress> inetSocketAddressList) {
        Selector Selector = new ConsistentHashLoadSelector(inetSocketAddressList, 128);
        return Selector;
    }

    /**
     * @author wzy
     * @description: 一直性hash的具体算法
     * @date 2023/10/1 12:16
     */
    private static class ConsistentHashLoadSelector implements Selector {
        //hash环,存储服务器节点
        private SortedMap<Integer, InetSocketAddress> circle = new TreeMap<>();
        //虚拟节点的个数
        private int virtualNodes;

        public ConsistentHashLoadSelector(List<InetSocketAddress> list, int virtualNodes) {
            //我们要尝试将节点转化为虚拟节点,进行挂载
            this.virtualNodes = virtualNodes;
            for (InetSocketAddress address : list) {
                //将每一个节点加入到circle
                addNodeToCircle(address);
            }
        }

        /**
         * @description: 将每个节点挂载到hash环上
         * @author wzy
         * @date 2023/10/1 12:35
         */
        private void addNodeToCircle(InetSocketAddress address) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(address.toString() + "-" + i);
                circle.put(hash, address);
            }
        }

        private void removeNodeFromCircle(InetSocketAddress address) {
            for (int i = 0; i < virtualNodes; i++) {
                int hash = hash(address.toString() + "-" + i);
                circle.remove(hash, address);
            }
        }

        @Override
        public InetSocketAddress getNext() {
            //hash环已经建立好了,接下来要对请求的要素做处理
            //根据请求的一些特征 选择服务器
            CrpcRequest crpcRequest = CRPC_REQUEST_THREAD_LOCAL.get();
            String requestId = String.valueOf(crpcRequest.getRequestId());
            //请求id做hash,字符串默认hash不太好
            int hash = hash(requestId);
            //判断该hash是否能直接落在服务器上,是否和服务器的hash一样
            if (!circle.containsKey(hash)) {
                //寻找离我最近的节点
                SortedMap<Integer, InetSocketAddress> tailMap = circle.tailMap(hash);
                hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
            }
            return circle.get(hash);
        }

        /**
         * @return
         * @description: 具体的hash算法
         * @params:
         * @author wzy
         * @date 2023/10/1 12:53
         */
        private int hash(String s) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");

            } catch (Exception e) {
                throw new RuntimeException();
            }
            byte[] digest = md.digest(s.getBytes(StandardCharsets.UTF_8));
            int res = 0;
            for (int i = 0; i < 4; i++) {
                int temp = digest[i] << ((3 - i) * 8);
                res = res | temp;
            }
            return res;
        }
    }

}
